"""
Uncertainty Weighting Strategy
不确定性权重策略实现

Reference: Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics
"""

from typing import Dict, List, Optional
import numpy as np
from ml_lib.core import Tensor
from ml_lib.nn.module import Parameter
from .base import WeightingStrategy


class UncertaintyWeighting(WeightingStrategy):
    """
    不确定性权重策略
    
    通过学习每个任务的不确定性参数来自适应调整任务权重。
    """
    
    def __init__(self, task_names: List[str], **kwargs):
        """
        初始化不确定性权重策略
        
        Args:
            task_names: 任务名称列表
        """
        super().__init__(task_names, **kwargs)
        
        # 初始化不确定性参数（对数方差）
        self.log_vars = {
            name: Parameter(Tensor([0.0], requires_grad=True)) 
            for name in task_names
        }
        
    def compute_weights(self, losses: Dict[str, Tensor], 
                       gradients: Optional[Dict[str, List[np.ndarray]]] = None,
                       **kwargs) -> Dict[str, float]:
        """
        计算任务权重
        
        Args:
            losses: 各任务的损失值字典
            gradients: 各任务的梯度字典（未使用）
            **kwargs: 其他参数
            
        Returns:
            任务权重字典
        """
        weights = {}
        
        for name in self.task_names:
            if name in losses:
                # 计算权重：1 / (2 * sigma^2)
                # 其中 sigma^2 = exp(log_var)
                log_var = self.log_vars[name].data.data[0]
                weight = 1.0 / (2.0 * np.exp(log_var))
                weights[name] = float(weight)
            else:
                weights[name] = 1.0
        
        return weights
    
    def compute_regularization_loss(self) -> Tensor:
        """
        计算正则化损失（不确定性项）
        
        Returns:
            正则化损失
        """
        reg_loss = Tensor([0.0])
        
        for name in self.task_names:
            # 正则化项：0.5 * log(sigma^2) = 0.5 * log_var
            log_var = self.log_vars[name].data
            reg_loss = reg_loss + log_var * 0.5
        
        return reg_loss
    
    def get_learnable_params(self) -> List[Parameter]:
        """获取可学习的不确定性参数"""
        return list(self.log_vars.values())
    
    def get_uncertainties(self) -> Dict[str, float]:
        """获取各任务的不确定性（标准差）"""
        uncertainties = {}
        for name in self.task_names:
            log_var = self.log_vars[name].data.data[0]
            sigma = np.sqrt(np.exp(log_var))
            uncertainties[name] = float(sigma)
        return uncertainties
    
    def reset(self):
        """重置权重策略状态"""
        super().reset()
        for name in self.task_names:
            self.log_vars[name].data.data = np.array([0.0])
    
    def __repr__(self):
        return f"UncertaintyWeighting(num_tasks={self.num_tasks})"